#pragma once

#include "znx_single_flight.h"
#include "znx_worker_proc.h"
#include "znx_session_pool.h"
#include "znx_udp_proc.h"
#include "znx_tcp_proc.h"
#include "comm/znx_slru.h"
#include "comm/znx_array.h"
#include "comm/znx_crc.h"


#define ZNX_RESOLVER_CACHE_SLOT     16

#define ZNX_RESOLVE_A               1
#define ZNX_RESOLVE_CNAME           5
#define ZNX_RESOLVE_PTR             12
#define ZNX_RESOLVE_MX              15
#define ZNX_RESOLVE_TXT             16
#define ZNX_RESOLVE_AAAA            28
#define ZNX_RESOLVE_SRV             33
#define ZNX_RESOLVE_DNAME           39

#define ZNX_RESOLVE_FORMERR         1
#define ZNX_RESOLVE_SERVFAIL        2
#define ZNX_RESOLVE_NXDOMAIN        3
#define ZNX_RESOLVE_NOTIMP          4
#define ZNX_RESOLVE_REFUSED         5
#define ZNX_RESOLVE_NO_NAMES        6
#define ZNX_RESOLVE_TIMEDOUT        ETIMEDOUT
#define ZNX_RESOLVE_INVALID_QUERY   1000
#define ZNX_RESOLVE_UNKNOW_ERR      1001
#define ZNX_RESOLVE_TRUNC_ERR       1002
#define ZNX_RESOLVE_MISMATCH_ERR    1003
#define ZNX_RESOLVE_BAD_QR          1004
#define ZNX_RESOLVE_INVALID_NAME    1005
#define ZNX_RESOLVE_UNSUPPORT_TYPE  1006


#define ZXN_RESOLVER_UDP_CONN_IDEL_MS   86400000 // 1day
#define ZNX_RESOLVER_TCP_CONN_IDEL_MS   5000     // 5s


typedef struct znx_resolver_s       znx_resolver_t;

typedef struct {
    znx_address_t                   address;
    znx_address_t                   bind_address;
    char                            address_text[ZNX_SOCKADDR_SIZE];
    char                            bind_address_text[ZNX_SOCKADDR_SIZE];
} znx_resolver_nameserver_t;


typedef struct {
    int                             eno;
    uint32_t                        crc;
    znx_str_t                       key;

#if 0
    int                             section;
    char                            *name;
    int                             type;
    int                             class;
#endif

    int                             ttl;
    
    znx_array_t                     *arr;
    znx_atomic_int64_t              refer;
    znx_cleanup_list_t              cleanup_list;
    int64_t                         expire_time;
} znx_resolver_result_t;


static inline znx_resolver_result_t *
znx_resolver_result_create()
{
    znx_resolver_result_t *res;
    res = znx_thread_calloc(sizeof(znx_resolver_result_t));
    res->arr = znx_array_create(8, sizeof(znx_str_t));
    res->refer = 1;
    znx_cleanup_list_init(&res->cleanup_list);
    return res;
}


static inline void
znx_resolver_result_hold(znx_resolver_result_t *res)
{
    int64_t old_refer;
    old_refer = znx_atomic_int64_fetch_add(&res->refer, 1);
    if (old_refer < 1) {
        ZNX_DEF_LOG_FATAL("old_refer must be greate than 0");
    }
}


static inline void
znx_resolver_result_hold_handler(void *data, znx_bool_t is_master)
{
    if (is_master) {
        return;
    }

    znx_resolver_result_t *result = data;
    znx_resolver_result_hold(result);
}


static inline void
znx_resolver_result_free(void *ret)
{
    znx_resolver_result_t *res =
        (znx_resolver_result_t *)ret;

    int64_t new_refer;
    new_refer = znx_atomic_int64_add_fetch(&res->refer, -1);
    if (new_refer < 0) {
        ZNX_DEF_LOG_FATAL("new_refer must be greate than or equal to 0");
    }

    if (new_refer != 0) {
        return;
    }

    znx_cleanup_list_exec(&res->cleanup_list);
    znx_array_destroy(res->arr);
    znx_thread_free(res);
}


static inline int
znx_resolver_result_cmp(void *a, void *b)
{
    znx_resolver_result_t *r_a = (znx_resolver_result_t *)a;
    znx_resolver_result_t *r_b = (znx_resolver_result_t *)b;

    if (r_a->crc > r_b->crc) {
        return 1;
    }

    if (r_a->crc < r_b->crc) {
        return -1;
    }

    if (r_a->key.len > r_b->key.len) {
        return 1;
    }

    if (r_a->key.len < r_b->key.len) {
        return -1;
    }

    return znx_strncmp(r_a->key.data,
        r_b->key.data, r_a->key.len);
}


typedef struct {
    znx_conn_t                      *conn;
    znx_resolver_nameserver_t       *nameserver;

    // buf = query - 2, used to save qlen.
    u_char                          *query;
    size_t                          qlen;

    // for tcp connect.
    size_t                          send_len;
    // the first two bytes are not included.
    size_t                          recv_len;
    size_t                          expect_data_len;
    u_char                          recv_len_buf[2];
    u_char                          *recv_data_buf;

    znx_resolver_result_t           *result;

    unsigned                        tcp_query:1;
    unsigned                        err:1;
} znx_resolver_task_ctx_t;


static inline znx_resolver_task_ctx_t *
znx_resolver_task_ctx_create()
{
    znx_resolver_task_ctx_t *ctx;
    ctx = znx_thread_calloc(sizeof(znx_resolver_task_ctx_t));
    return ctx;
}


static inline void
znx_resolver_task_ctx_reset(znx_resolver_task_ctx_t *ctx)
{
    // reuse ctx->query.
    if (ctx->conn) {
        size_t  tag_len = 0;
        tag_len += 128;
        tag_len += ZNX_SOCKADDR_SIZE;
        tag_len += ZNX_SOCKADDR_SIZE;
        znx_g_buf_grow(tag_len);

        u_char *end = znx_snprintf(g_buf, tag_len, "DNS-%d-%s-%s",
            ctx->tcp_query,
            ctx->nameserver->address_text,
            ctx->nameserver->bind_address_text);

        znx_str_t tag = {
            .data = g_buf,
            .len = (size_t)(end - g_buf)
        };

        znx_event_loop_remove_conn(ctx->conn);

        if (ctx->err) {
            znx_conn_close(ctx->conn);
            znx_thread_free(ctx->conn);
        } else {
            znx_g_session_pool_put(&tag, ctx->conn,
                ctx->tcp_query ?
                    ZNX_RESOLVER_TCP_CONN_IDEL_MS : ZXN_RESOLVER_UDP_CONN_IDEL_MS,
                NULL);
        }

        ctx->conn = NULL;
    }

    ctx->nameserver = NULL;

    ctx->send_len = 0;
    ctx->recv_len = 0;
    ctx->expect_data_len = 0;
    ctx->recv_len_buf[0] = (u_char)(0);
    ctx->recv_len_buf[1] = (u_char)(0);

    if (ctx->recv_data_buf) {
        znx_thread_free(ctx->recv_data_buf);
        ctx->recv_data_buf = NULL;
    }

    ctx->tcp_query = 0;
    ctx->err = 0;
}


static inline void
znx_resolver_task_ctx_destroy(znx_resolver_task_ctx_t *ctx)
{
    znx_resolver_task_ctx_reset(ctx);

    if (ctx->query) {
        // for tcp query.
        // the first two bytes are used to store the qlen.
        znx_thread_free(ctx->query - 2);
    }

    znx_thread_free(ctx);
}


typedef struct znx_resolver_task_s  znx_resolver_task_t;

typedef void (*znx_resolver_task_handle_pt)(znx_resolver_task_t *task);

struct znx_resolver_task_s {
    // request ctx.
    void                            *data;
    znx_resolver_task_handle_pt     handler;

    const char                      *tag;
    int                             qtype;
    znx_str_t                       query;
    unsigned                        tcp_query:1;
    unsigned                        force_update:1;
    int64_t                         timeout_ms;

    /* use in single_flight */
    znx_dispatch_event_t            dispatch_event;

    uint32_t                        crc;
    znx_str_t                       key;

    size_t                          step;
    znx_resolver_task_ctx_t         *ctx;
    int                             last_eno;

    znx_resolver_t                  *owner_resolver;
    znx_resolver_result_t           *result;
};


void znx_resolver_task_done(znx_dispatch_event_t *dispatch_event);

static inline znx_resolver_task_t *
znx_resolver_task_create(void *data,
    znx_resolver_task_handle_pt handler,
    const char *tag, int qtype, znx_str_t *query,
    znx_bool_t tcp_query, znx_bool_t force_update, int64_t timeout_ms)
{
    znx_resolver_task_t *task =
        znx_thread_calloc(sizeof(znx_resolver_task_t));
    task->data = data;
    task->handler = handler;
    task->dispatch_event.data = task;
    task->dispatch_event.handler = znx_resolver_task_done;
    task->tag = tag;
    task->qtype = qtype;
    task->query = *query;

    if (tcp_query) {
        task->tcp_query = 1;
    }

    if (force_update) {
        task->force_update = 1;
    }

    if (timeout_ms < 60) {
        timeout_ms = 60;
    }

    task->timeout_ms = timeout_ms;

    size_t key_len = 0;
    key_len += strlen(tag);
    key_len += query->len;
    key_len += 128;

    u_char *key_buf =
        znx_thread_malloc(key_len);
    u_char *key_end =
        znx_snprintf(key_buf, key_len, "%s-%d-%v", tag, qtype, query);
    key_len = (size_t)(key_end - key_buf);

    task->crc = znx_crc32_long(key_buf, key_len);
    task->key.data = key_buf;
    task->key.len = key_len;

    return task;
}


static inline void
znx_resolver_task_destroy(znx_resolver_task_t *task)
{
    if (task->ctx != NULL) {
        znx_resolver_task_ctx_destroy(task->ctx);
    }

    if (task->result != NULL) {
        znx_resolver_result_free(task->result);
    }

    if (task->key.data) {
        znx_thread_free(task->key.data);
    }

    znx_thread_free(task);
}


static inline int
znx_resolver_task_cmp(znx_dispatch_event_t *a,
    znx_dispatch_event_t *b)
{
    znx_resolver_task_t *task_a =
        znx_queue_data(a, znx_resolver_task_t, dispatch_event);
    znx_resolver_task_t *task_b =
        znx_queue_data(b, znx_resolver_task_t, dispatch_event);

    if (task_a->crc > task_b->crc) {
        return 1;
    }

    if (task_a->crc < task_b->crc) {
        return -1;
    }

    if (task_a->key.len > task_b->key.len) {
        return 1;
    }

    if (task_a->key.len < task_b->key.len) {
        return -1;
    }

    return znx_strncmp(task_a->key.data,
        task_b->key.data, task_a->key.len);
}


struct znx_resolver_s {
    znx_array_t                     *names;
    znx_bool_t                      has_cache;
    znx_spinlock_t                  locks[ZNX_RESOLVER_CACHE_SLOT];
    znx_slru_t                      *slrus[ZNX_RESOLVER_CACHE_SLOT];
    znx_single_flight_t             sf;
    znx_cleanup_list_t              cleanup;
};


// the names elements is znx_resolver_nameserver_t.
// we will make a copy of the string.
// if names is NULL, we will load /etc/resolv.conf nameservers.
// if resolver create failed, will return NULL.
znx_resolver_t *znx_resolver_create(znx_array_t *nameservers,
    int64_t max_cache_item);


void znx_resolver_destroy(znx_resolver_t *resolver);


void znx_resolver_query(znx_resolver_t *resolver,
    znx_resolver_task_t *task);


// ------------------------- for testing --------------------------

znx_bool_t znx_testing_resolver_build_name_query(znx_resolver_task_t *task);

int znx_testing_resolver_decode_name(u_char *buf, size_t n, size_t *posp);

